import ast
from PIL import Image, ImageDraw
import onnxruntime as ort
import torchvision.transforms as transforms


if __name__ == '__main__':
    # 打开图像并调整大小
    image_source = Image.open("D:\\Workspace\\deal\\ada0c42e8bc94f559be0fe0a74644535.jpg").convert("RGB")

    print("原始的图片的类型", type(image_source))

    height, width = 640, 640

    image_source = image_source.resize((height, width))  # 根据模型要求调整大小




    # 定义图像转换
    preprocess = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 替换为你的模型要求的均值和标准差
    ])

    # 应用转换
    image_train = preprocess(image_source)

    print("预测的图片的类型", type(image_train))

    # 增加一个维度，变为 (1, C, H, W)
    image_train = image_train.unsqueeze(0)



    # 加载 ONNX 模型
    session = ort.InferenceSession("D:\\Workspace\\deal\\yolov5s.onnx", providers=["CUDAExecutionProvider", "CPUExecutionProvider"])

    # 运行模型进行推理，指定输出和输入
    inputs = session.get_inputs()
    outputs = session.get_outputs()

    print("inputs :", type(inputs), len(inputs), inputs[0])
    print("outputs:", type(outputs), len(outputs), outputs[0])

    input_name = session.get_inputs()[0].name
    output_name = session.get_outputs()[0].name
    # 初步学习使用 run 来预测图片上的目标
    outputs = session.run([output_name], {input_name: image_train.numpy()})

    print("预测结果：")
    print("outputs len:", len(outputs))
    print(outputs)

    print()

    targets = outputs[0][0]
    print("预测到的目标数量为：", len(targets))

    # 创建一个绘图对象
    draw = ImageDraw.Draw(image_source)
    # 置信度的过滤线，阈值
    obj_conf = 0.25

    names_dict = ast.literal_eval(session.get_modelmeta().custom_metadata_map.get("names"))

    for target in targets:
        # 每一个 target 就是识别到的一个目标，一个方框。

        target_box_conf = target[4]

        # 第一次筛选
        if target_box_conf < obj_conf:
            continue

        # 多个类别的置信度中的最大置信度的值
        target_max_score_value = 0
        # 多个类别的置信度中的最大置信度的类别的索引
        target_max_score_index = 0
        # 类别的置信度数据是从 target 数组的索引位置 5 的地方开始
        temp_index = 5

        while temp_index < len(target):

            score = target[temp_index]

            # 寻找最大的置信度值及类别索引
            if score > target_max_score_value:
                target_max_score_value = score
                target_max_score_index = temp_index

            temp_index += 1

        if target_max_score_index < 6:
            continue

        target_max_score_index = target_max_score_index - 5

        # 第二次筛选
        if target_max_score_value * target_box_conf < obj_conf:
            continue

        # box_center_x, box_center_y, box_width, box_height 分别是目标中心的位置和宽高
        box_center_x, box_center_y, box_width, box_height = int(target[0]), int(target[1]), int(target[2]), int(target[3])

        # 将 box_center_x, box_center_y, box_width, box_height 转为 (x1, y1) (x2, y2) 坐标
        # 定义方框的位置（左上角和右下角坐标）
        temp_top_left = (box_center_x - (box_width / 2), box_center_y - (box_height / 2))
        temp_bottom_right = (box_center_x + (box_width / 2), box_center_y + (box_height / 2))

        # 在图像上绘制方框
        draw.rectangle([temp_top_left, temp_bottom_right], outline='blue', width=1)

        name = names_dict.get(target_max_score_index)

        print("检测到时目标：", name, temp_top_left, temp_bottom_right)

    # 保存或显示图像
    image_source.show()  # 显示图像





